import numpy as np
import matplotlib.pyplot as plt
import cv2
from sklearn.linear_model import LinearRegression


def plot_image_and_r2(image, x1, x2, y, k1, k2, b, r2):
    # 将 BGR 格式转换为 RGB 格式
    image = cv2.cvtColor(image.copy(), cv2.COLOR_BGR2RGB)

    # 创建一个图形和两个子图
    fig, ax1 = plt.subplots(1, 1, figsize=(10, 5))

    # 设置窗口标题
    fig.canvas.manager.set_window_title("建模结果")

    # 在第一个子图中显示图片
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title('Box')

    # 在第二个子图中显示拟合直线
    ax2 = fig.add_subplot(122)
    ax2.scatter(x1, y, color='b', label='X1 vs Y')
    ax2.scatter(x2, y, color='g', label='X2 vs Y')

    # 绘制拟合直线
    x1_line = np.linspace(min(x1), max(x1), 100)
    x2_line = np.linspace(min(x2), max(x2), 100)
    y_line = k1 * x1_line + k2 * x2_line + b

    ax2.plot(x1_line, y_line, color='r', label='Fitted Line')
    ax2.set_title('Fitted Line (k1={:.5f}, k2={:.5f}, b={:.5f}, r2={:.5f})'.format(k1, k2, b, r2))
    ax2.legend()

    # 显示图形
    plt.tight_layout()
    plt.show()


# 示例数据
np.random.seed(0)
X1 = np.random.rand(100) * 10  # 生成100个随机的X1值
X2 = np.random.rand(100) * 10  # 生成100个随机的X2值
y = 2 * X1 + 3 * X2 + 5 + np.random.randn(100) * 2  # 生成对应的y值

# 多元线性回归
X = np.column_stack((X1, X2))
reg = LinearRegression().fit(X, y)
k1, k2 = reg.coef_
b = reg.intercept_

# 计算 R^2 值
y_pred = reg.predict(X)
r2 = reg.score(X, y)

# 示例图片
image = np.zeros((100, 100, 3), dtype=np.uint8)

# 调用函数绘制图像和拟合直线
plot_image_and_r2(image, X1, X2, y, k1, k2, b, r2)


